# -*- coding: utf-8 -*-
"""
@Time    : 2024/10/22 11:23 
@Author  : ZhangShenao 
@File    : indexer.py 
@Desc    : 索引器

从指定位置加载文档,并将文档切分成chunks分块
并调用Embedding模型,生成嵌入向量,存储至向量数据库
"""
import os

from langchain_community.document_loaders import PDFPlumberLoader, TextLoader, \
    UnstructuredWordDocumentLoader, UnstructuredPowerPointLoader, UnstructuredExcelLoader, CSVLoader, \
    UnstructuredMarkdownLoader, UnstructuredXMLLoader, UnstructuredHTMLLoader
from langchain_core.documents import Document
from langchain_core.vectorstores import VectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 定义文档解析加载器字典
# 根据文档的扩展名类型选择对应的文档解析加载器类和输入参数
DOCUMENT_LOADER_MAPPING = {
    ".pdf": (PDFPlumberLoader, {}),
    ".txt": (TextLoader, {"encoding": "utf8"}),
    ".doc": (UnstructuredWordDocumentLoader, {}),
    ".docx": (UnstructuredWordDocumentLoader, {}),
    ".ppt": (UnstructuredPowerPointLoader, {}),
    ".pptx": (UnstructuredPowerPointLoader, {}),
    ".xlsx": (UnstructuredExcelLoader, {}),
    ".csv": (CSVLoader, {}),
    ".md": (UnstructuredMarkdownLoader, {}),
    ".xml": (UnstructuredXMLLoader, {}),
    ".html": (UnstructuredHTMLLoader, {}),
}


def indexing(file_path: str, vector_store: VectorStore) -> list[Document]:
    """
    Indexing 索引过程
    :param file_path: 文件路径
    :param vector_store: 向量数据库
    :return 文档分块列表
    """

    print("Indexing开始执行!")

    # 加载文档
    docs = load(file_path=file_path)

    # 文档分割
    chunks = split(docs=docs, chunk_size=512, chunk_overlap=128)

    # Embedding存储
    embedding_and_store(docs=chunks, vector_store=vector_store)

    # 返回文档分块列表
    print("Indexing执行完成")
    return chunks


def load(file_path: str) -> list[Document]:
    """
    加载文档
    :param file_path: PDF文件路径
    :return: 加载后的文档列表
    """

    # 解析文件扩展名类型
    ext = os.path.splitext(file_path)[1]
    # 根据文件扩展名类型,获取解析器类和参数
    loader_mapping = DOCUMENT_LOADER_MAPPING.get(ext)
    if loader_mapping is None:
        raise ValueError(f"不支持的文件类型: {ext}")

    # 解包元组,获取文档解析加载器类和参数
    loader_class, loader_args = loader_mapping
    # 创建文档解析加载器实例,并传入文档文件路径和扩展参数
    loader = loader_class(file_path, **loader_args)

    # 加载文档对象,并返回
    docs = loader.load()
    print("文档加载完成!")
    return docs


def split(docs: list[Document], chunk_size: int, chunk_overlap: int) -> list[Document]:
    """
    对文档进行分割,返回chunks文档块列表
    :param docs: 原始文档列表
    :param chunk_size: 文档分块大小
    :param chunk_overlap: 分块间重叠部分大小
    :return: 文档分块列表
    """

    # 创建RecursiveCharacterTextSplitter递归文档分割器
    # 通过层次化的分隔符(从双换行符到单字符)拆分文本
    # 旨在保持文本的结构和连贯性
    # 优先考虑自然边界如段落和句子
    # 文档分割的三大关键要素:
    #     - chunk_size: 分块大小,用于确定每个分块的长度
    #     - chunk_overlap: 分块重叠大小,控制在相邻分块间的重叠字符数,可以控制分块间上下文的连贯性
    #     - separators: 分隔符列表,用于确定不同分块的边界
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)

    # 分割文档,返回文档块列表
    chunks = splitter.split_documents(docs)
    print(f"文档块分割完成! 分块数量: {len(chunks)}")
    return chunks


def embedding_and_store(docs: list[Document], vector_store: VectorStore) -> None:
    """
    将对文档块进行Embedding,生成向量,并存入向量数据库
    :param docs: 文档列表
    :param vector_store: 向量数据库
    """

    # 将文档块存入向量数据库
    doc_ids = vector_store.add_documents(docs)
    print(f"向量数据库存储完成! doc_id 数量: {len(doc_ids)}")
